/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#ifndef SRC_FLOW_FLOWMANAGER_H_
#define SRC_FLOW_FLOWMANAGER_H_

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <fstream>
#include <limits>
#include "Flow.h"
#include "Protocol.h"
#include "FlowCache.h"
#include "Cache.h"
#include "OutputManager.h"
#include "FlowTable.h"
#include "FlowSearchOptions.h"
#include "json.hpp"

namespace aiengine {

class FlowManager {
public:
	explicit FlowManager(const std::string &name);
	explicit FlowManager(): FlowManager("FlowManager") {}
	virtual ~FlowManager();

	static const int flowTimeRefreshRate = 64;
	static const int flowTimeout = 180; // Seconds

	void setReleaseFlows(bool value) { release_flows_ = value; }
	bool haveReleaseFlows() const { return release_flows_; }

	void add(const SharedPointer<Flow> &flow);
	void remove(const SharedPointer<Flow> &flow);
	SharedPointer<Flow>& find(unsigned long hash1, unsigned long hash2);
	
	// This method will not update the position of the network flow internally
	// This member function is only used by the stacks for show the flow to the user
	SharedPointer<Flow> find(const FlowSearchOptions &fsos);
	void updateTimers(const std::time_t current_time);

	void setFlowCache(FlowCachePtr cache) { flow_cache_ = cache; }
	void setTCPInfoCache(Cache<TCPInfo>::CachePtr cache) { tcp_info_cache_ = cache; }

	void setTimeout(int timeout) { timeout_ = timeout; }
	int getTimeout() const { return timeout_; }
	int getTotalFlows() const { return (int)flowTable_.size();}

	int32_t getTotalProcessFlows() const { return total_process_flows_;}
	int32_t getTotalTimeoutFlows() const { return total_timeout_flows_;}

	void showFlows() const { showFlows(std::numeric_limits<int>::max()); }
	void showFlows(int limit) const ;
	void showFlows(const std::string &protoname) const { showFlows(protoname, std::numeric_limits<int>::max());}
	void showFlows(const std::string &protoname, int limit) const;

	void showFlows(std::basic_ostream<char> &out, int limit) const;
	void showFlows(std::basic_ostream<char> &out, int limit, const std::string &protoname) const;
	void showFlows(Json &out, int limit) const;
	void showFlows(Json &out, int limit, const std::string &protoname) const;
	void showFlows(std::basic_ostream<char> &out, std::function<bool (const Flow&)> condition) const { show_flows(out, condition); };
	void showFlows(Json &out, std::function<bool (const Flow&)> condition) const { show_flows(out, condition); };

	int64_t getAllocatedMemory() const;

	// Method for flush the flows of the FlowManager in order to analyse more times
	// This method just retrieve the flows to their corresponding caches and frees the
	// memory attached to them
	void flush();
	// Also can be flushed flows depending the L7 protocol
	void flush(const std::string &name);

	// Close all the flows that their timeout expired
	void purge();
	void purge(std::time_t current_time);

#if defined(STAND_ALONE_TEST) || defined(TESTING)
	void showFlowsByTime();
	Cache<TCPInfo>::CachePtr getTCPInfoCache() const { return tcp_info_cache_; }
#endif

	void statistics(std::basic_ostream<char> &out) const { out << *this;}
        void statistics() const { statistics(OutputManager::getInstance()->out());}

	friend std::ostream& operator<< (std::ostream &out, const FlowManager &fm);

	FlowTable getFlowTable() const { return flowTable_;}
	SharedPointer<Flow> getLastProcessFlow() const { return lookup_flow_; }

#if defined(PYTHON_BINDING)
	// Methods for exposing the class to python iterable methods
	auto begin() { return flowTable_.begin(); }
	auto end() { return flowTable_.end(); }
#endif
	void setProtocol(ProtocolPtrWeak proto) { protocol_ = proto; }


private:
	void show_flows(std::basic_ostream<char> &out, std::function<bool (const Flow&)> condition) const;
	void show_flows(Json &out, std::function<bool (const Flow&)> condition) const;
	void print_pretty_flow(std::basic_ostream<char> &out, const Flow &flow, std::time_t current_time) const;
	void print_pretty_flow(Json &out, const Flow &flow) const;
	void release(const SharedPointer<Flow> &flow);

	std::string name_;
	int32_t total_process_flows_ = 0;
	int32_t total_timeout_flows_ = 0;
	int timeout_ = flowTimeout;
	bool release_flows_ = false;
	FlowTable flowTable_;
	FlowCachePtr flow_cache_ = nullptr;
	Cache<TCPInfo>::CachePtr tcp_info_cache_ = nullptr;
	ProtocolPtrWeak protocol_;
	SharedPointer<Flow> lookup_flow_ = nullptr; // cacheable flow;
#if defined(PYTHON_BINDING)
	std::ofstream output_term_;
#endif
	boost::posix_time::ptime data_time_ = boost::posix_time::microsec_clock::local_time();
};

typedef std::shared_ptr<FlowManager> FlowManagerPtr;
typedef std::weak_ptr<FlowManager> FlowManagerPtrWeak;

} // namespace aiengine

#endif  // SRC_FLOW_FLOWMANAGER_H_
